import os
import zipfile
import datetime
import argparse
from pathlib import Path

def backup_files(source_dir, backup_dir, include_ext=None, exclude_dirs=None):
    """
    备份文件到压缩包
    :param source_dir: 源目录
    :param backup_dir: 备份存储目录
    :param include_ext: 只备份指定扩展名，如'.txt,.docx'
    :param exclude_dirs: 排除的目录，如'node_modules,.git'
    """
    try:
        # 创建备份目录（如果不存在）
        Path(backup_dir).mkdir(parents=True, exist_ok=True)
        
        # 生成备份文件名（包含时间戳）
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        backup_filename = f"backup_{timestamp}.zip"
        backup_path = os.path.join(backup_dir, backup_filename)
        
        # 解析过滤条件
        include_extensions = [ext.lower() for ext in include_ext.split(',')] if include_ext else None
        exclude_directories = [d.lower() for d in exclude_dirs.split(',')] if exclude_dirs else []
        
        # 创建压缩文件
        with zipfile.ZipFile(backup_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
            file_count = 0
            # 遍历源目录
            for root, dirs, files in os.walk(source_dir):
                # 过滤目录
                dirs[:] = [d for d in dirs if d.lower() not in exclude_directories]
                
                for file in files:
                    # 过滤文件扩展名
                    if include_extensions:
                        file_ext = os.path.splitext(file)[1].lower()
                        if file_ext not in include_extensions:
                            continue
                    
                    file_path = os.path.join(root, file)
                    # 计算相对路径，保持目录结构
                    arcname = os.path.relpath(file_path, source_dir)
                    zipf.write(file_path, arcname)
                    file_count += 1
        
        print(f"备份完成！共备份 {file_count} 个文件")
        print(f"备份文件: {backup_path}")
        print(f"文件大小: {os.path.getsize(backup_path) / 1024:.2f} KB")
        
    except Exception as e:
        print(f"备份失败: {str(e)}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='文件自动备份脚本')
    parser.add_argument('source', help='需要备份的源目录')
    parser.add_argument('backup', help='备份文件存储目录')
    parser.add_argument('--include', help='只备份指定扩展名，用逗号分隔，如.txt,.docx')
    parser.add_argument('--exclude', help='排除的目录，用逗号分隔，如.node_modules,.git')
    
    args = parser.parse_args()
    
    backup_files(
        source_dir=args.source,
        backup_dir=args.backup,
        include_ext=args.include,
        exclude_dirs=args.exclude
    )
